`include "define.sv"

module fifo_to_array (
    input clk,
    input reset,
    input [1:0] matrix_type,       // 矩阵类型
    // 来自上游的输入
    input [31:0]  data_in_a,        // 单个元素输出（最大32位）
    input         data_valid_in_a,  // 输出数据有效信号
    input [31:0]  data_in_b,        // 单个元素输出（最大32位）
    input         b_trans_done,  // 输出数据有效信号
    input         fifo_b_data_available,  // FIFO B数据有效信号
    output        reg data_req_b,       // B矩阵转置请求信号
    // 下游信号
    output reg [31:0] data_out_a [0:31],  // 32个通道的输出数据（非压缩数组）
    output reg wr_en_a [0:31],            // 32个通道的写使能
    output reg [31:0] data_out_b [0:31],  // 32个通道的输出数据（非压缩数组）
    output reg wr_en_b [0:31],            // 32个通道的写使能
    input array_ready,              // 阵列空闲信号
    output reg done                // 写入一轮完成标志
);

// ========== 新增FIFO控制状态机 ==========
typedef enum logic [1:0] {
    IDLE,
    WAIT_COMPLETE,
    READING
} fifo_state_t;

fifo_state_t fifo_b_state;

always @(posedge clk or posedge reset) begin
    if (reset || !array_ready) begin
        fifo_b_state <= IDLE;
        data_req_b <= 0;
    end else begin
        case (fifo_b_state)
            IDLE: begin
                if (b_trans_done) begin
                    fifo_b_state <= WAIT_COMPLETE;
                end
            end
            
            WAIT_COMPLETE: begin
                data_req_b <= 1'b1;  // 开始连续读取
                fifo_b_state <= READING;
            end
            
            READING: begin
                if (b_trans_done || !fifo_b_data_available) begin
                    data_req_b <= 1'b0;  // 停止读取
                    fifo_b_state <= IDLE;
                end
            end
        endcase
    end
end

// 根据矩阵类型选择FIFO数量
reg [4:0] max_fifo_num_a;
always @(*) begin
    case (matrix_type)
        `m8n32k16:  max_fifo_num_a = 5'd7;   // 8 FIFOs (0-7)
        `m16n16k16: max_fifo_num_a = 5'd15;  // 16 FIFOs (0-15)
        `m32n8k16:  max_fifo_num_a = 5'd31;  // 32 FIFOs (0-31)
        default:    max_fifo_num_a = 5'd15;  // 默认16 FIFOs
    endcase
end

reg [4:0] max_fifo_num_b;
always @(*) begin
    case (matrix_type)
        `m8n32k16:  max_fifo_num_b = 5'd31;  // 32 FIFOs (0-31)
        `m16n16k16: max_fifo_num_b = 5'd15;  // 16 FIFOs (0-15)
        `m32n8k16:  max_fifo_num_b = 5'd7;   // 8 FIFOs (0-7)
        default:    max_fifo_num_b = 5'd15;  // 默认16 FIFOs
    endcase
end

// 内部状态寄存器
reg [4:0] current_fifo_a, current_fifo_b;  // 当前填充的FIFO索引
reg [4:0] count_a, count_b;                // 当前FIFO已填充数据计数
reg a_done, b_done;                        // 通道A/B完成标志

integer i; // 循环变量

always @(posedge clk or posedge reset) begin
    if (reset) begin
        data_req_b <= 1'b0;
    end else if(b_trans_done)begin
        data_req_b <= 1'b1;
    end else begin
        data_req_b <= 1'b0;
    end
end

always @(posedge clk) begin
    if (reset || !array_ready) begin
        // 复位或下游未准备好时初始化状态
        current_fifo_a <= 0;
        current_fifo_b <= 0;
        count_a <= 0;
        count_b <= 0;
        a_done <= 0;
        b_done <= 0;
        done <= 0;
        // 所有写使能置0
        for (i = 0; i < 32; i = i + 1) begin
            wr_en_a[i] <= 0;
            wr_en_b[i] <= 0;
        end
    end else begin
        // 默认写使能置0
        for (i = 0; i < 32; i = i + 1) begin
            wr_en_a[i] <= 0;
            wr_en_b[i] <= 0;
        end

        // 处理通道A数据
        if (!a_done && data_valid_in_a) begin
            // 写入数据到当前FIFO
            data_out_a[current_fifo_a] <= data_in_a;
            wr_en_a[current_fifo_a] <= 1'b1;
            
            // 更新计数器
            if (count_a == 15) begin
                count_a <= 0;
                if (current_fifo_a == max_fifo_num_a) begin
                    a_done <= 1'b1;  // A通道所有FIFO填满
                end else begin
                    current_fifo_a <= current_fifo_a + 1;  // 切换到下一个FIFO
                end
            end else begin
                count_a <= count_a + 1;
            end
        end

        // B通道处理（使用fifo_b_rd_en作为有效信号）
        if (data_req_b && fifo_b_data_available && !b_done) begin
            data_out_b[current_fifo_b] <= data_in_b;
            wr_en_b[current_fifo_b] <= 1'b1;
            
            // 更新计数器
            if (count_b == 15) begin
                count_b <= 0;
                if (current_fifo_b == max_fifo_num_b) begin
                    b_done <= 1'b1;
                end else begin
                    current_fifo_b <= current_fifo_b + 1;
                end
            end else begin
                count_b <= count_b + 1;
            end
        end


        // 检查是否完成
        done <= a_done && b_done;
    end
end

endmodule